{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import math\n",
    "import pickle\n",
    "\n",
    "from paws.data_to_mp4 import encode_tensor_to_video, decode_video_to_tensor\n",
    "\n",
    "from paws.model_to_grid import make_source_2d, class_grid_to_medium_grid, make_sensor_2d, simulation_2d, point_cloud_voxelization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:\\replica\\apartment_1\\habitat\\mesh_semantic.ply\n",
      "C:\\replica\\apartment_1\\habitat\\info_semantic.json\n"
     ]
    }
   ],
   "source": [
    "dir_path = \"C:\\\\replica\\\\apartment_1\"\n",
    "\n",
    "\n",
    "ply_path = os.path.join(dir_path,\"habitat\",\"mesh_semantic.ply\")\n",
    "print(ply_path)\n",
    "json_path = os.path.join(dir_path,\"habitat\",\"info_semantic.json\")\n",
    "print(json_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------------------------------------------------+\n",
      "|                  kspaceFirstOrder-CUDA v1.3                   |\n",
      "+---------------------------------------------------------------+\n",
      "| Reading simulation configuration:                        Done |\n",
      "| Selected GPU device id:                                     0 |\n",
      "| GPU device name:                NVIDIA GeForce RTX 4070 SUPER |\n",
      "| Number of CPU threads:                                      1 |\n",
      "| Processor name:         13th Gen Intel(R) Core(TM) i5-13600KF |\n",
      "+---------------------------------------------------------------+\n",
      "|                      Simulation details                       |\n",
      "+---------------------------------------------------------------+\n",
      "| Domain dimensions:                                  256 x 256 |\n",
      "| Medium type:                                               2D |\n",
      "| Simulation time steps:                                   1207 |\n",
      "+---------------------------------------------------------------+\n",
      "|                        Initialization                         |\n",
      "+---------------------------------------------------------------+\n",
      "| Memory allocation:                                       Done |\n",
      "| Data loading:                                            Done |\n",
      "| Elapsed time:                                           0.00s |\n",
      "+---------------------------------------------------------------+\n",
      "| FFT plans creation:                                      Done |\n",
      "| Pre-processing phase:                                    Done |\n",
      "| Elapsed time:                                           1.29s |\n",
      "+---------------------------------------------------------------+\n",
      "|                    Computational resources                    |\n",
      "+---------------------------------------------------------------+\n",
      "| Current host memory in use:                             351MB |\n",
      "| Current device memory in use:                          1296MB |\n",
      "| Expected output file size:                              301MB |\n",
      "+---------------------------------------------------------------+\n",
      "|                          Simulation                           |\n",
      "+----------+----------------+--------------+--------------------+\n",
      "| Progress |  Elapsed time  |  Time to go  |  Est. finish time  |\n",
      "+----------+----------------+--------------+--------------------+\n",
      "|     0%   |        0.001s  |      0.603s  |  10/06/24 15:57:11 |\n",
      "|     5%   |        0.020s  |      0.369s  |  10/06/24 15:57:11 |\n",
      "|    10%   |        0.037s  |      0.329s  |  10/06/24 15:57:11 |\n",
      "|    15%   |        0.055s  |      0.308s  |  10/06/24 15:57:11 |\n",
      "|    20%   |        0.074s  |      0.294s  |  10/06/24 15:57:11 |\n",
      "|    25%   |        0.091s  |      0.271s  |  10/06/24 15:57:11 |\n",
      "|    30%   |        0.108s  |      0.250s  |  10/06/24 15:57:11 |\n",
      "|    35%   |        0.124s  |      0.229s  |  10/06/24 15:57:11 |\n",
      "|    40%   |        0.141s  |      0.211s  |  10/06/24 15:57:11 |\n",
      "|    45%   |        0.159s  |      0.193s  |  10/06/24 15:57:11 |\n",
      "|    50%   |        0.175s  |      0.174s  |  10/06/24 15:57:11 |\n",
      "|    55%   |        0.197s  |      0.161s  |  10/06/24 15:57:11 |\n",
      "|    60%   |        0.215s  |      0.142s  |  10/06/24 15:57:12 |\n",
      "|    65%   |        0.233s  |      0.125s  |  10/06/24 15:57:12 |\n",
      "|    70%   |        0.252s  |      0.108s  |  10/06/24 15:57:12 |\n",
      "|    75%   |        0.269s  |      0.089s  |  10/06/24 15:57:12 |\n",
      "|    80%   |        0.287s  |      0.071s  |  10/06/24 15:57:12 |\n",
      "|    85%   |        0.308s  |      0.054s  |  10/06/24 15:57:12 |\n",
      "|    90%   |        0.327s  |      0.036s  |  10/06/24 15:57:12 |\n",
      "|    95%   |        0.343s  |      0.018s  |  10/06/24 15:57:12 |\n",
      "+----------+----------------+--------------+--------------------+\n",
      "| Elapsed time:                                           0.36s |\n",
      "+---------------------------------------------------------------+\n",
      "| Sampled data post-processing:                            Done |\n",
      "| Elapsed time:                                           0.00s |\n",
      "+---------------------------------------------------------------+\n",
      "|                            Summary                            |\n",
      "+---------------------------------------------------------------+\n",
      "| Peak host memory in use:                                351MB |\n",
      "| Peak device memory in use:                             1296MB |\n",
      "+---------------------------------------------------------------+\n",
      "| Total execution time:                                   1.83s |\n",
      "+---------------------------------------------------------------+\n",
      "|                       End of computation                      |\n",
      "+---------------------------------------------------------------+\n",
      "Write video to c:\\Users\\Administrator\\Desktop\\codes\\PAWS\\data\\test_output_pressure.mp4\n",
      "[5.0000000e+00 4.7164102e+00 3.9235163e+00 ... 8.1596074e-05 8.0922720e-05\n",
      " 8.0986829e-05]\n",
      "Write out meta to c:\\Users\\Administrator\\Desktop\\codes\\PAWS\\data\\test_output_pressure.pkl\n"
     ]
    }
   ],
   "source": [
    "\n",
    "Nx = 256\n",
    "Ny = 256\n",
    "\n",
    "dx = 2e-2\n",
    "dy = 2e-2\n",
    "\n",
    "\n",
    "\n",
    "cls_grid = np.zeros([Nx,Ny],dtype=int)\n",
    "\n",
    "\n",
    "medium = class_grid_to_medium_grid(cls_grid,{})\n",
    "\n",
    "source = make_source_2d(128,128,2,Nx,Ny,5)\n",
    "\n",
    "sensor = make_sensor_2d(np.ones([Nx,Ny],dtype=bool))\n",
    "\n",
    "\n",
    "\n",
    "sensor_data = simulation_2d(Nx,Ny,dx,dy,source,medium,sensor)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# save path of project\n",
    "base_dir = os.getcwd()\n",
    "save_dir = os.path.join(base_dir,\"data\")\n",
    "\n",
    "save_base = os.path.basename(dir_path)\n",
    "\n",
    "\n",
    "print(sensor_data)\n",
    "\n",
    "pressure_dist = sensor_data[\"p\"].copy()\n",
    "pressure_dist = pressure_dist.reshape((pressure_dist.shape[0], Nx, Ny))\n",
    "\n",
    "save_filename = \"test_output\"\n",
    "\n",
    "save_path = os.path.join(save_dir,save_filename)\n",
    "\n",
    "#save pressure distribution\n",
    "video_path = save_path + \"_pressure.mp4\"\n",
    "meta_path = save_path + \"_pressure.pkl\"\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "v_scales = encode_tensor_to_video(\n",
    "    tensor=pressure_dist, \n",
    "    video_path=video_path,\n",
    ")\n",
    "# v_mins: (T,)\n",
    "# v_maxs: (T,)\n",
    "###\n",
    "\n",
    "\n",
    "print(v_scales)\n",
    "\n",
    "meta = {\n",
    "    \"v_scales\": v_scales\n",
    "}\n",
    "pickle.dump(meta, open(meta_path, \"wb\"))\n",
    "print(\"Write out meta to {}\".format(meta_path))\n",
    "\n",
    "#save medium data\n",
    "medium_save_pth = save_path + \"_medium.npy\"\n",
    "medium_data = {\"sound_speed\":medium.sound_speed,\n",
    "                \"density\":medium.density}\n",
    "\n",
    "np.save(medium_save_pth,medium_data)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "k-wave",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
